package com.fanshuai.spring.boot.xa.mysql;

import com.fanshuai.spring.boot.sql.StatementCreator;
import com.mysql.cj.jdbc.MysqlXid;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import javax.sql.XAConnection;
import javax.sql.XADataSource;
import javax.transaction.xa.XAResource;
import javax.transaction.xa.Xid;
import java.nio.charset.StandardCharsets;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.List;
import java.util.concurrent.atomic.LongAdder;

@Service
@Slf4j
public class CustomXAService {
    @Resource
    private XADataSource xaDataSource1;
    @Resource
    private XADataSource xaDataSource2;

    private LongAdder longAdder = new LongAdder();

    /**
     * XA事务处理  自定义事务管理器
     */
    public void executeXA(StatementCreator creator1, StatementCreator creator2) {
        try {
            XAConnection xaConnection1 = xaDataSource1.getXAConnection();
            XAConnection xaConnection2 = xaDataSource2.getXAConnection();

            List<PreparedStatement> statements1 = creator1.createPreparedStatement(xaConnection1.getConnection());
            List<PreparedStatement> statements2 = creator2.createPreparedStatement(xaConnection2.getConnection());
            executeXA0(xaConnection1, xaConnection2, statements1, statements2);
        } catch (Exception e) {
            log.error("get connection error, ", e);
        }
    }

    /**
     * XA事务处理  使用两个mysql实例连接  自定义事务管理器  资源管理器：resource1, resource2
     * @param statements1
     * @param statements2
     */
    private void executeXA0(XAConnection xaConnection1, XAConnection xaConnection2,
                            List<PreparedStatement> statements1, List<PreparedStatement> statements2) throws SQLException {
        XAResource resource1 = xaConnection1.getXAResource();
        XAResource resource2 = xaConnection2.getXAResource();

        byte[] gtrid = ("g" + getId()).getBytes(StandardCharsets.UTF_8); //全局事务id
        int format = 1;
        byte[] bqual1 = ("b" + getId()).getBytes(StandardCharsets.UTF_8); //mysql本地事务1
        byte[] bqual2 = ("b" + getId()).getBytes(StandardCharsets.UTF_8); //mysql本地事务2
        Xid xid1 = new MysqlXid(gtrid, bqual1, format);
        Xid xid2 = new MysqlXid(gtrid, bqual2, format);

        try {
            //rm1的事务分支
            resource1.start(xid1, XAResource.TMNOFLAGS);  //One of TMNOFLAGS, TMJOIN, or TMRESUME
            for (PreparedStatement statement : statements1) {
                statement.execute();
            }
            resource1.end(xid1, XAResource.TMSUCCESS);

            //rm2的事务分支
            resource2.start(xid2, XAResource.TMNOFLAGS);  //One of TMNOFLAGS, TMJOIN, or TMRESUME
            for (PreparedStatement statement : statements2) {
                statement.execute();
            }
            resource2.end(xid2, XAResource.TMSUCCESS);

            //prepare
            int res1 = resource1.prepare(xid1);
            int res2 = resource2.prepare(xid2);

            //commit/rollback
            if (res1 == XAResource.XA_OK && res2 == XAResource.XA_OK) {
                resource1.commit(xid1, false);
                resource2.commit(xid2, false);
            } else {
                resource1.rollback(xid1);
                resource2.rollback(xid2);
            }
        } catch (Exception e) {
            log.error("xa error, ", e);
            try {
                //rollback
                resource1.rollback(xid1);
                resource2.rollback(xid2);
            } catch (Exception ex) {
                log.error("xa rollback error, ", ex);
            }
        } finally {
            try {
                for (Statement statement : statements1) {
                    statement.close();
                }
                for (Statement statement : statements2) {
                    statement.close();
                }
                xaConnection1.close();
                xaConnection2.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    private String getId() {
        long v = longAdder.longValue();
        longAdder.add(1);

        return String.valueOf(v % 100000);
    }
}
